#include "exif_writer.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>

// EXIF条目结构
typedef struct {
    uint16_t tag;
    uint16_t type;
    uint32_t count;
    uint32_t value_offset;
} exif_entry_t;

// IFD结构
typedef struct {
    exif_entry_t* entries;
    int entry_count;
    int entry_capacity;
} exif_ifd_data_t;

// EXIF上下文结构
struct exif_context {
    exif_ifd_data_t ifds[3]; // IFD0, ExifIFD, GPS IFD
    char* make;
    char* model;
    double latitude;
    double longitude;
    uint16_t iso;
    int has_gps;
    int has_iso;
};

// 创建EXIF上下文
exif_context_t* exif_create_context(void) {
    exif_context_t* ctx = (exif_context_t*)malloc(sizeof(exif_context_t));
    if (!ctx) return NULL;
    
    memset(ctx, 0, sizeof(exif_context_t));
    
    // 初始化IFD数组
    for (int i = 0; i < 3; i++) {
        ctx->ifds[i].entries = NULL;
        ctx->ifds[i].entry_count = 0;
        ctx->ifds[i].entry_capacity = 0;
    }
    
    return ctx;
}

// 销毁EXIF上下文
void exif_destroy_context(exif_context_t* ctx) {
    if (!ctx) return;
    
    for (int i = 0; i < 3; i++) {
        free(ctx->ifds[i].entries);
    }
    
    free(ctx->make);
    free(ctx->model);
    free(ctx);
}

// 添加条目到IFD
static int add_entry_to_ifd(exif_ifd_data_t* ifd, uint16_t tag, uint16_t type, 
                           uint32_t count, uint32_t value_offset) {
    if (ifd->entry_count >= ifd->entry_capacity) {
        int new_capacity = ifd->entry_capacity == 0 ? 10 : ifd->entry_capacity * 2;
        exif_entry_t* new_entries = realloc(ifd->entries, new_capacity * sizeof(exif_entry_t));
        if (!new_entries) return -1;
        ifd->entries = new_entries;
        ifd->entry_capacity = new_capacity;
    }
    
    ifd->entries[ifd->entry_count].tag = tag;
    ifd->entries[ifd->entry_count].type = type;
    ifd->entries[ifd->entry_count].count = count;
    ifd->entries[ifd->entry_count].value_offset = value_offset;
    ifd->entry_count++;
    
    return 0;
}

// 添加字符串
int exif_add_string(exif_context_t* ctx, exif_ifd_t ifd, uint16_t tag, const char* value) {
    if (!ctx || !value) return -1;
    
    if (tag == EXIF_TAG_MAKE) {
        free(ctx->make);
        ctx->make = (char*)malloc(strlen(value) + 1);
        if (!ctx->make) return -1;
        memcpy(ctx->make, value, strlen(value) + 1);
        return 0;
    } else if (tag == EXIF_TAG_MODEL) {
        free(ctx->model);
        ctx->model = (char*)malloc(strlen(value) + 1);
        if (!ctx->model) return -1;
        memcpy(ctx->model, value, strlen(value) + 1);
        return 0;
    }
    
    // 对于其他字符串，添加到对应的IFD
    return add_entry_to_ifd(&ctx->ifds[ifd], tag, EXIF_TYPE_ASCII, strlen(value) + 1, 0);
}

// 添加短整型
int exif_add_short(exif_context_t* ctx, exif_ifd_t ifd, uint16_t tag, uint16_t value) {
    if (!ctx) return -1;
    
    if (tag == EXIF_TAG_ISO_SPEED_RATINGS) {
        ctx->iso = value;
        ctx->has_iso = 1;
        return 0;
    }
    
    return add_entry_to_ifd(&ctx->ifds[ifd], tag, EXIF_TYPE_SHORT, 1, value);
}

// 添加长整型
int exif_add_long(exif_context_t* ctx, exif_ifd_t ifd, uint16_t tag, uint32_t value) {
    if (!ctx) return -1;
    return add_entry_to_ifd(&ctx->ifds[ifd], tag, EXIF_TYPE_LONG, 1, value);
}

// 添加有理数
int exif_add_rational(exif_context_t* ctx, exif_ifd_t ifd, uint16_t tag, 
                     uint32_t numerator, uint32_t denominator) {
    (void)numerator; (void)denominator; // 当前未用，保留接口
    if (!ctx) return -1;
    return add_entry_to_ifd(&ctx->ifds[ifd], tag, EXIF_TYPE_RATIONAL, 1, 0);
}

// 添加GPS坐标
int exif_add_gps_coordinates(exif_context_t* ctx, double latitude, double longitude) {
    if (!ctx) return -1;
    
    ctx->latitude = latitude;
    ctx->longitude = longitude;
    ctx->has_gps = 1;
    return 0;
}

// 写入大端16位
static void write_be16(FILE *f, uint16_t v) {
    uint8_t b[2] = { (uint8_t)((v >> 8) & 0xFF), (uint8_t)(v & 0xFF) };
    fwrite(b, 1, 2, f);
}

// 写入小端16位到缓冲区
static void append_le16(uint8_t *buf, size_t *idx, uint16_t v) {
    buf[(*idx)++] = (uint8_t)(v & 0xFF);
    buf[(*idx)++] = (uint8_t)((v >> 8) & 0xFF);
}

// 写入小端32位到缓冲区
static void append_le32(uint8_t *buf, size_t *idx, uint32_t v) {
    buf[(*idx)++] = (uint8_t)(v & 0xFF);
    buf[(*idx)++] = (uint8_t)((v >> 8) & 0xFF);
    buf[(*idx)++] = (uint8_t)((v >> 16) & 0xFF);
    buf[(*idx)++] = (uint8_t)((v >> 24) & 0xFF);
}

// 写入字节到缓冲区
static void append_bytes(uint8_t *buf, size_t *idx, const void *src, size_t len) {
    memcpy(buf + *idx, src, len);
    *idx += len;
}

// 构建EXIF载荷
static size_t build_exif_payload(uint8_t *payload, exif_context_t* ctx) {
    size_t idx = 0;
    
    // "Exif\0\0"
    const uint8_t exif_id[6] = { 'E','x','i','f',0,0 };
    append_bytes(payload, &idx, exif_id, 6);

    // TIFF头 (小端序)
    append_bytes(payload, &idx, (const uint8_t*)"II", 2);
    append_le16(payload, &idx, 0x002A);
    append_le32(payload, &idx, 8);

    const uint32_t tiff_start = 6;
    const uint32_t ifd0_offset = 8;
    
    // 计算数据布局
    uint32_t data_offset = ifd0_offset + 2 + ctx->ifds[0].entry_count * 12 + 4;
    
    // 添加Make和Model字符串的偏移
    if (ctx->make) {
        data_offset += strlen(ctx->make) + 1;
    }
    if (ctx->model) {
        data_offset += strlen(ctx->model) + 1;
    }
    
    // 添加ExifIFD指针
    if (ctx->has_iso) {
        data_offset += 2 + 1 * 12 + 4; // ExifIFD大小
    }
    
    // 添加GPS IFD指针
    if (ctx->has_gps) {
        data_offset += 2 + 5 * 12 + 4; // GPS IFD大小
        data_offset += 24 + 24; // GPS数据
    }
    
    // 写入IFD0
    size_t expected_ifd0_start = tiff_start + ifd0_offset;
    if (idx < expected_ifd0_start) {
        size_t pad = expected_ifd0_start - idx;
        memset(payload + idx, 0, pad);
        idx += pad;
    }
    
    // 计算IFD0条目数
    int ifd0_entries = ctx->ifds[0].entry_count;
    if (ctx->make) ifd0_entries++;
    if (ctx->model) ifd0_entries++;
    if (ctx->has_iso) ifd0_entries++;
    if (ctx->has_gps) ifd0_entries++;
    
    append_le16(payload, &idx, ifd0_entries);
    
    uint32_t current_offset = ifd0_offset + 2 + ifd0_entries * 12 + 4;
    
    // 写入Make条目
    if (ctx->make) {
        append_le16(payload, &idx, EXIF_TAG_MAKE);
        append_le16(payload, &idx, EXIF_TYPE_ASCII);
        append_le32(payload, &idx, strlen(ctx->make) + 1);
        append_le32(payload, &idx, current_offset);
        current_offset += strlen(ctx->make) + 1;
    }
    
    // 写入Model条目
    if (ctx->model) {
        append_le16(payload, &idx, EXIF_TAG_MODEL);
        append_le16(payload, &idx, EXIF_TYPE_ASCII);
        append_le32(payload, &idx, strlen(ctx->model) + 1);
        append_le32(payload, &idx, current_offset);
        current_offset += strlen(ctx->model) + 1;
    }
    
    // 写入ExifIFD指针
    if (ctx->has_iso) {
        append_le16(payload, &idx, 0x8769); // ExifIFD指针
        append_le16(payload, &idx, EXIF_TYPE_LONG);
        append_le32(payload, &idx, 1);
        append_le32(payload, &idx, current_offset);
        current_offset += 2 + 1 * 12 + 4;
    }
    
    // 写入GPS IFD指针
    if (ctx->has_gps) {
        append_le16(payload, &idx, 0x8825); // GPS IFD指针
        append_le16(payload, &idx, EXIF_TYPE_LONG);
        append_le32(payload, &idx, 1);
        append_le32(payload, &idx, current_offset);
        current_offset += 2 + 5 * 12 + 4 + 24 + 24;
    }
    
    // 写入其他IFD0条目
    for (int i = 0; i < ctx->ifds[0].entry_count; i++) {
        append_le16(payload, &idx, ctx->ifds[0].entries[i].tag);
        append_le16(payload, &idx, ctx->ifds[0].entries[i].type);
        append_le32(payload, &idx, ctx->ifds[0].entries[i].count);
        append_le32(payload, &idx, ctx->ifds[0].entries[i].value_offset);
    }
    
    // Next IFD offset
    append_le32(payload, &idx, 0);
    
    // 写入字符串数据
    if (ctx->make) {
        append_bytes(payload, &idx, ctx->make, strlen(ctx->make) + 1);
    }
    if (ctx->model) {
        append_bytes(payload, &idx, ctx->model, strlen(ctx->model) + 1);
    }
    
    // 写入ExifIFD
    if (ctx->has_iso) {
        append_le16(payload, &idx, 1); // 条目数
        append_le16(payload, &idx, EXIF_TAG_ISO_SPEED_RATINGS);
        append_le16(payload, &idx, EXIF_TYPE_SHORT);
        append_le32(payload, &idx, 1);
        append_le16(payload, &idx, ctx->iso);
        append_le16(payload, &idx, 0);
        append_le32(payload, &idx, 0); // next IFD
    }
    
    // 写入GPS IFD
    if (ctx->has_gps) {
        append_le16(payload, &idx, 5); // 条目数
        
        // GPSVersionID
        append_le16(payload, &idx, EXIF_TAG_GPS_VERSION_ID);
        append_le16(payload, &idx, EXIF_TYPE_BYTE);
        append_le32(payload, &idx, 4);
        append_bytes(payload, &idx, (uint8_t[4]){2,3,0,0}, 4);
        
        // LatitudeRef
        char lat_ref = ctx->latitude >= 0.0 ? 'N' : 'S';
        append_le16(payload, &idx, EXIF_TAG_GPS_LATITUDE_REF);
        append_le16(payload, &idx, EXIF_TYPE_ASCII);
        append_le32(payload, &idx, 2);
        append_bytes(payload, &idx, (uint8_t[4]){(uint8_t)lat_ref, 0, 0, 0}, 4);
        
        // Latitude
        append_le16(payload, &idx, EXIF_TAG_GPS_LATITUDE);
        append_le16(payload, &idx, EXIF_TYPE_RATIONAL);
        append_le32(payload, &idx, 3);
        append_le32(payload, &idx, current_offset - 24 - 24);
        
        // LongitudeRef
        char lon_ref = ctx->longitude >= 0.0 ? 'E' : 'W';
        append_le16(payload, &idx, EXIF_TAG_GPS_LONGITUDE_REF);
        append_le16(payload, &idx, EXIF_TYPE_ASCII);
        append_le32(payload, &idx, 2);
        append_bytes(payload, &idx, (uint8_t[4]){(uint8_t)lon_ref, 0, 0, 0}, 4);
        
        // Longitude
        append_le16(payload, &idx, EXIF_TAG_GPS_LONGITUDE);
        append_le16(payload, &idx, EXIF_TYPE_RATIONAL);
        append_le32(payload, &idx, 3);
        append_le32(payload, &idx, current_offset - 24);
        
        append_le32(payload, &idx, 0); // next IFD
        
        // 写入GPS数据
        double alat = fabs(ctx->latitude);
        double alon = fabs(ctx->longitude);
        
        uint32_t lat_deg = (uint32_t)floor(alat);
        double lat_min_f = (alat - lat_deg) * 60.0;
        uint32_t lat_min = (uint32_t)floor(lat_min_f);
        double lat_sec_f = (lat_min_f - lat_min) * 60.0;
        
        uint32_t lon_deg = (uint32_t)floor(alon);
        double lon_min_f = (alon - lon_deg) * 60.0;
        uint32_t lon_min = (uint32_t)floor(lon_min_f);
        double lon_sec_f = (lon_min_f - lon_min) * 60.0;
        
        // Latitude RATIONAL[3]
        append_le32(payload, &idx, lat_deg);
        append_le32(payload, &idx, 1);
        append_le32(payload, &idx, lat_min);
        append_le32(payload, &idx, 1);
        append_le32(payload, &idx, (uint32_t)llround(lat_sec_f * 10000.0));
        append_le32(payload, &idx, 10000);
        
        // Longitude RATIONAL[3]
        append_le32(payload, &idx, lon_deg);
        append_le32(payload, &idx, 1);
        append_le32(payload, &idx, lon_min);
        append_le32(payload, &idx, 1);
        append_le32(payload, &idx, (uint32_t)llround(lon_sec_f * 10000.0));
        append_le32(payload, &idx, 10000);
    }
    
    return idx;
}

// 写入EXIF到文件
int exif_write_to_file(exif_context_t* ctx, const char* input_file, const char* output_file) {
    if (!ctx || !input_file || !output_file) return -1;
    
    FILE *in = fopen(input_file, "rb");
    if (!in) return -1;
    
    FILE *out = fopen(output_file, "wb");
    if (!out) {
        fclose(in);
        return -1;
    }
    
    // 读取并验证SOI
    uint8_t soi_bytes[2];
    if (fread(soi_bytes, 1, 2, in) != 2) {
        fclose(in);
        fclose(out);
        return -1;
    }
    
    uint16_t soi = ((uint16_t)soi_bytes[0] << 8) | soi_bytes[1];
    if (soi != 0xFFD8) {
        fclose(in);
        fclose(out);
        return -1;
    }
    
    // 写入SOI
    fwrite(soi_bytes, 1, 2, out);
    
    // 构建并写入EXIF载荷
    uint8_t payload[4096];
    size_t payload_size = build_exif_payload(payload, ctx);
    
    write_be16(out, 0xFFE1); // APP1
    uint16_t app1_length = (uint16_t)(payload_size + 2);
    write_be16(out, app1_length);
    fwrite(payload, 1, payload_size, out);
    
    // 复制剩余数据
    uint8_t buf[4096];
    size_t n;
    while ((n = fread(buf, 1, sizeof(buf), in)) > 0) {
        fwrite(buf, 1, n, out);
    }
    
    fclose(in);
    fclose(out);
    return 0;
}

// 示例函数：添加EXIF数据
void add_exif_data(const char* filename) {
    exif_context_t* ctx = exif_create_context();
    if (!ctx) return;
    
    // 添加基本EXIF信息
    exif_add_string(ctx, EXIF_IFD0, EXIF_TAG_MAKE, "Cannice");
    exif_add_string(ctx, EXIF_IFD0, EXIF_TAG_MODEL, "EOS 5D Mark IV");
    exif_add_short(ctx, EXIF_EXIF_IFD, EXIF_TAG_ISO_SPEED_RATINGS, 100);
    exif_add_gps_coordinates(ctx, 39.9042, 116.4074);
    
    // 写入文件
    char output_file[256];
    snprintf(output_file, sizeof(output_file), "%s_with_exif.jpg", filename);
    exif_write_to_file(ctx, filename, output_file);
    
    exif_destroy_context(ctx);
} 